#!/usr/bin/env python3
#coding=utf-8

import pymongo  # cmongo/golang:1.12已经安装了pymongo, etcd
import etcd
import string
import sys
import os
import time
import threading
import subprocess
import json
from bson import json_util
import random
import logging
import coloredlogs

# Note:
#   此测试严重依赖于docker集群环境

# oss ip within docker env
# (NOTE phoenixxliu): 当在本地环境测试时，可用如下命令加以修正：
#   cp docker/oss-test.py a.py && sed -i 's/30\./31\./g' a.py &&  python3 a.py 3.2
ETCD_HOST='172.30.0.13'
MASTER_HOST='172.30.0.14'
MASTER_PORT='8888'
MONGOS_HOSTS=['172.30.0.16', '172.30.0.17', '172.30.0.18']
CONFIG_HOSTS=['172.30.0.19', '172.30.0.20', '172.30.0.21']
MONGOD_HOSTS=['172.30.0.22', '172.30.0.23', '172.30.0.24', '172.30.0.25', '172.30.0.26', '172.30.0.27']

POPULATE_NUM=100
# 可以在本地配置LOCAL_PYTEST变量，这样某些动作可以忽略
local=False

logger = logging.getLogger('main')
synclogger = logging.getLogger('cluster_sync')
poplogger = logging.getLogger('populate')
failoverlogger = logging.getLogger('failover')

logging.getLogger("etcd").setLevel(logging.WARNING)
logging.getLogger("urllib3").setLevel(logging.WARNING)

# Initialize coloredlogs.
coloredlogs.install(level='DEBUG')

def get_value_from_env(k, default):
    v = os.getenv(k)
    if v is not None:
        return v
    else:
        return default

NAME_PREFIX=get_value_from_env('NAME_PREFIX', 'cmongo')

SRC_CLUSTER_3_2=NAME_PREFIX + '-cluster-3.2'
DEST_CLUSTER_3_2=NAME_PREFIX + '-cluster-3.2_2'
SRC_CLUSTER_3_6=NAME_PREFIX + '-cluster-3.6'
DEST_CLUSTER_3_6=NAME_PREFIX + '-cluster-3.6_2'
SRC_CLUSTER_4_0=NAME_PREFIX + '-cluster-4.0'
DEST_CLUSTER_4_0=NAME_PREFIX + '-cluster-4.0_2'

cluster_list = [SRC_CLUSTER_3_2, SRC_CLUSTER_3_6, SRC_CLUSTER_4_0, DEST_CLUSTER_3_2, DEST_CLUSTER_3_6, DEST_CLUSTER_4_0]

salt_string=''
version=''

etcd_client = etcd.Client(host=ETCD_HOST, port=2379)

def get_salt_string():
    root_user_salt = ''
    md5_salt = ''
    for l in open('utils/salt_string.go'):
        if 'ROOT_USER_SALT' in l:
            root_user_salt = l.strip().replace(' ', '').split('=')[1][1:-1]
        elif 'MD5_SALT' in l:
            md5_salt = l.strip().replace(' ', '').split('=')[1][1:-1]

    return root_user_salt + '-' + md5_salt

def get_client(host, port, user='', passwd=''):
    user_passwd=''
    if user != '':
        user_passwd = user + ':' + passwd + '@'
    #mongodb_url = 'mongodb://' + user_passwd + host + ':' + str(port) + '/' + dbname + '?authSource=admin'
    mongodb_url = 'mongodb://' + user_passwd + host + ':' + str(port) + '/' +  'admin'
    client = pymongo.MongoClient(mongodb_url)
    dblist = client.list_database_names()
    #print('dbs on %s:%d :%s, client addr: %s' %(host, port, dblist, client.address))
    return client

def get_collection(client, dbname, colname):
    db = client[dbname]

    col = db[colname]
    #collist = db.list_collection_names()
    return col

def raw_cmd(cmd_str, src_cluster='', dst_cluster='', rs_name='', ip='', port='', id=''):
    cmd_str = cmd_str.replace('$MASTER_HOST', MASTER_HOST)
    cmd_str = cmd_str.replace('$MASTER_PORT', MASTER_PORT)
    cmd_str = cmd_str.replace('$ETCD_HOST', ETCD_HOST)
    cmd_str = cmd_str.replace('$CLUSTER', src_cluster)
    cmd_str = cmd_str.replace('$SRC_CLUSTER', src_cluster)
    cmd_str = cmd_str.replace('$DEST_CLUSTER', dst_cluster)
    cmd_str = cmd_str.replace('$RS_NAME', rs_name)
    cmd_str = cmd_str.replace('$IP', ip)
    cmd_str = cmd_str.replace('$PORT', port)
    cmd_str = cmd_str.replace('$ID', id)
    return cmd_str

'''
def print_in_color(str):
    print("\033[34m%s\033[0m" % str)

def print_failure_in_color(str):
    print("\033[47;31m%s\033[0m" % str)
'''

def failure(str, exit_code=-1):
    logger.critical(str)
    sys.exit(exit_code)

def wait_until_etcd_key_not_existed(key, timeout_s):
    i=0
    ok=False
    while i < timeout_s:
        try:
            etcd_client.read(key)
            if i%10 == 0:
                logger.warning('%s still exists on etcd after %ds, continue to sleep' %(key, i))
            time.sleep(1)
        except etcd.EtcdKeyNotFound:
            ok = True
            break
        i += 1

    return ok

def assert_set_cluster_sync(col, expected, cluster):
    # 查询目标集群是否有了数据
    i = 0
    n = 0
    while i < 18:
        try:
            n = col.find().count()
            # 同步后目标集群的记录数 应等于 源集群的记录数
            assert n == expected
            break
        except Exception as e:
            synclogger.info('i=%d -> %s count:%d, expected: %d, sleeping 10s ...' %(i, col.full_name, n, expected))
            synclogger.debug('etcd sync status: %s' %(etcd_client.get('/cluster_sync/' + cluster + '/job').value))
            time.sleep(10)
            i += 1

    if n != expected:
        synclogger.error("FAILURE: actual: %d, expected: %d records within %s" %(n, expected, col.full_name))
        return False
    else:
        synclogger.info('%s counts ok' %(col.full_name))
        return True

# 检查cluster sync后，主实例mongouser用户及密码是否同步到目标集群
def assert_user_transferred(host, port, src_cluster):
    conn_str = 'mongodb://mongouser:' + src_cluster + '@' + host + ':' + str(port) + '/admin'
    c = pymongo.MongoClient(conn_str, serverSelectionTimeoutMS=3000)

    try:
        c.admin.command('ismaster')
    except Exception as e:
        failure('FAILURE: cannot login with mongouser to cluster: %s, err: %s' %(src_cluster, e))

def print_config_databases(host, port, user, passwd):
    c = get_client(host, port, user, passwd)
    t = get_collection(c, 'config', 'databases')
    while True:
        synclogger.debug('config.databases info:')
        res=t.find()
        for r in res:
            print(r)
        time.sleep(5)

def assert_detach_cluster_sync(dest, src_cluster):
    cmd_set_cluster_sync='mastertest -cmd detach_cluster_sync -master_address $MASTER_HOST:$MASTER_PORT -src_region 1 -src_set 1 ' \
                         '-dest_region 1 -dest_set 1 -src_cluster $SRC_CLUSTER -dest_cluster $DEST_CLUSTER'

    cmd = raw_cmd(cmd_set_cluster_sync, get_value_from_env('SRC_CLUSTER', src_cluster),
                  get_value_from_env('DEST_CLUSTER', dest))
    synclogger.debug('CMD: %s' % cmd)

    if NAME_PREFIX != 'cmongo':
        # 在本地就不要执行detach了，免得中间状态查看不了
        # 已经打印出cmd，必要时复制执行
        return

    output = os.popen(cmd)
    synclogger.debug(output.read())

    k = '/cluster_sync/' + dest
    if not wait_until_etcd_key_not_existed(k, 90):
        failure('delete etcd key: %s failed' %(k))
    else:
        logger.info('detach cluster: %s sync ok' %dest)

def check_CLUSTER_SYNC(host, port, dest, passwd, src_cluster):
    synclogger.info("calling set_cluster_sync for %s" % src_cluster)
    client = get_client(host, port, dest, passwd)
    #client.read_preference
    cmd_set_cluster_sync='mastertest -cmd set_cluster_sync -master_address $MASTER_HOST:$MASTER_PORT -src_region 1 -src_set 1 ' \
                         '-dest_region 1 -dest_set 1 -src_cluster $SRC_CLUSTER -dest_cluster $DEST_CLUSTER'

    cmd = raw_cmd(cmd_set_cluster_sync, get_value_from_env('SRC_CLUSTER', src_cluster),
                  get_value_from_env('DEST_CLUSTER', dest))
    synclogger.debug('CMD: %s' % cmd)
    output = os.popen(cmd)
    synclogger.debug(output.read())

    # pr_cfg = threading.Thread(target=print_config_databases, args=(CONFIG_HOSTS[0], 7003, user, passwd)) ###################
    # pr_cfg.start()

    # 获得目标集群的目标collection实例, db1:t
    db1_t = get_collection(client, 'db1', 't')
    db2_t = get_collection(client, 'db2', 't')
    pytest_t = get_collection(client, 'pytest', 't')
    ok1 = assert_set_cluster_sync(db1_t, 100, dest)
    ok2 = assert_set_cluster_sync(db2_t, 200, dest)

    ok3 = assert_set_cluster_sync(pytest_t, POPULATE_NUM, dest)
    if not ok3:
        # ok3 如果都没成，那前两个很可能也没成。经过这段时间都没成，只能报告失败
        failure('db1.t db2.t pytset.t all failed')
    # ok3 成了，ok1,ok2也许会因等待时间不够而fail了，可以但检查一遍
    if not ok1:
        if not assert_set_cluster_sync(db1_t, 100, dest):
            failure('db1.t still failed', -2)
    if not ok2:
        if not assert_set_cluster_sync(db2_t, 200, dest):
            failure('db2.t still failed', -3)

    assert_user_transferred(host, port, src_cluster)

    assert_detach_cluster_sync(dest, src_cluster)
    client.close()

def check_SHOWDB(cluster):
    cmd_showdb = 'mastertest -cmd showdb -cluster $CLUSTER -master_address $MASTER_HOST:$MASTER_PORT'
    cmd = raw_cmd(cmd_showdb, cluster)
    output = os.popen(cmd).read()
    logger.debug('cmd: %s, output: %s' %(cmd,output))
    if 'db_list:\"db1\"' not in output:
        failure("FAILURE: showdb not contains the correct results")

def get_needed_container(cluster, etcd_url, is_mongod=True): # return containerName
    try:
        cluster_info = etcd_client.read('/cluster/' + cluster).value
        # 反序列化
        #print(cluster_info)
        text = json.loads(cluster_info)
        if is_mongod:
            rs = text['rs_list'][0]
            mongod_num = len(rs['mongod_list'])
            # get the last one
            container = rs['mongod_list'][mongod_num-1]
        else:
            # NOTE phoenix: 如果是WT3.2的集群，这里需要取'proxy_list'而非'mongos_list'
            mongos_num = len(text['mongos_list'])
            # get the last one
            container = text['mongos_list'][mongos_num-1]
        #print(container['container_name']) #debug
        return container['container_name']
    except etcd.EtcdKeyNotFound:
        logger.error('etcd key not found')

# container_name is like:"10.175.102.28:7015:12302018"
def split_container_name(container_name):
    list = container_name.split(':')
    if len(list) != 3:
        logger.error('unexpected container_name')
    return list[0], list[1], list[2]

def check_VOTE_NEW_PRIMARY(cluster, new_primary):
    rs_name = cluster+'_0'
    ip, port, id = split_container_name(new_primary)
    cmd_vote = 'mastertest -master_address $MASTER_HOST:$MASTER_PORT -cmd vote_new_primary -cluster $CLUSTER \
        -rsname $RS_NAME -ip $IP -port $PORT -id $ID -election_timeout_sec 10'
    cmd = raw_cmd(cmd_vote, cluster, '', rs_name,ip,port, id)
    output = os.popen(cmd).read()
    logger.debug('cmd: %s, output: %s' %(cmd,output))

    if 'vote new primary result' not in output:
        failure("FAILURE: vote_new_primary run failed, err: %s" %(output))
    # TODO：更合理的检测方式，比如rs.status()或者遍历etcd元数据

def check_DELETE_CONTAINER(cluster, container, is_mongod=True):
    rs_name = cluster+'_0'
    ip, port, id = split_container_name(container)
    cmd_delete = 'mastertest -master_address $MASTER_HOST:$MASTER_PORT -cmd delcontainer -cluster $CLUSTER \
        -rsname $RS_NAME -ip $IP -port $PORT -id $ID -container_type 0'
    if not is_mongod:
        cmd_delete = 'mastertest -master_address $MASTER_HOST:$MASTER_PORT -cmd delcontainer -cluster $CLUSTER \
            -rsname $RS_NAME -ip $IP -port $PORT -id $ID -container_type 1'
    cmd = raw_cmd(cmd_delete, cluster, '', rs_name,ip,port, id)
    output = os.popen(cmd).read()
    logger.debug('cmd: %s, output: %s' %(cmd,output))

    if 'DelContainer success done' not in output:
        failure("FAILURE: del_container run failed, error: %s" %(output))
    # TODO：更合理的检测方式

def check_DESTROY(cluster):
    cmd_destroy='mastertest -master_address $MASTER_HOST:$MASTER_PORT -cmd destroy -cluster $CLUSTER'
    cmd = raw_cmd(cmd_destroy, cluster)
    output = os.popen(cmd).read()
    logger.debug('cmd: %s, output: %s' %(cmd,output))
    # 这个关键字依赖于代码的日志，不提倡这样check，否则日志描述一变，这里就得变
    if 'DelCluster success' not in output:
        failure("mastertest del cluster failed, err: %v" %(output))

    if not wait_until_etcd_key_not_existed('/cluster/' + cluster, 60):
        failure('destroy cluster: %s failed' %cluster)
    else:
        logger.debug('destroy %s ok' %cluster)

def entrypoint(mongos_host, mongos_port, src, dest, passwd, etcd_host):
    check_SHOWDB(src)
    check_CLUSTER_SYNC(mongos_host, mongos_port, dest, passwd, src)

    container = get_needed_container(src, ETCD_HOST)
    # @phoenixxliu 当主从切换正在发生时，它容易失败, https://git.code.oa.com/cmongo/cmongo/issues/33
    if local == False:
        check_VOTE_NEW_PRIMARY(src, container)
    #check_DELETE_CONTAINER(src, container)
    # TODO： 删除proxy节点可能需要mock一下通知NETAGENT和l5的接口
    #container = get_needed_container(src, ETCD_HOST, False)
    #check_DELETE_CONTAINER(src, container, False)

    # 对于OCI，由于pytest是紧跟集群创建， 所以只能在4.0集群创建后执行destroy操作; 本地暂手工destroy
    # 尽量destroy目标集群，源集群可能还会干些事
    if version == '4.0' and NAME_PREFIX == 'cmongo':
        check_DESTROY(dest)

def populate(host, port, user, passwd):
    thread_id = threading.get_ident()
    cli = get_client(host, port, user, passwd)
    pytest_t = get_collection(cli, 'pytest', 't')
    i=1
    while i <= POPULATE_NUM:
        key = 'k' + str(i)
        try:
            pytest_t.update_one({key : i}, {'$set' : {'value' : str(thread_id) + '-' + str(i)} }, upsert=True)
            #print('\rkey: %s inserted\t' %(key), end='')
            if i%10 == 0:
                poplogger.info('%d records inserted' %(i))
        except Exception as e:
            poplogger.warning('update key: %s except: %s retrying ...' %(key, e))
            time.sleep(2)
            continue

        i += 1
        time.sleep(1)
    cli.close()
    poplogger.info('thread %d populate %d records into cluster: %s done' %(thread_id, POPULATE_NUM, user))

def get_primary(cluster):
    primary = ''
    try:
        cluster_info_str = etcd_client.get('/cluster/' + cluster).value
        #print(cluster_info_str)
        cluster_info = json.loads(cluster_info_str)
        mongod_list=cluster_info['rs_list'][0]["mongod_list"]
        mongod_containers =[]
        for e in mongod_list:
            for k, v in e.items():
                if k == 'container_name':
                    mongod_containers.append(v)
        #print('cluster: %s has mongod: %s' %(cluster, str(mongod_containers)))

        node1 = mongod_containers[0]
        mongo_shell_cmd='mongo mongodb://' + cluster + ':' + cluster + '-' + salt_string + \
                        '@' + node1[0:node1.rfind(':')] + '/admin --eval "db.isMaster()"'
        # 先去掉前面的2行输出
        res=os.popen(mongo_shell_cmd).read().split('\n', 2)[2]
        #print('res=', res)

        primary=os.popen('echo "' + res + '" | grep primary').read().split(':', 1)[1].strip()[:-1]

        # 由于mongo_shell_cmd返回的是BSON结构，json_util解析好像也有问题，所以换成了shell来找到primary
        # print('a=', type(a), json.loads(a))
        # primary=json.loads(json_util.dumps(a))

        #print('cluster %s primary: %s' %(cluster,primary))
    except Exception as e:
        logger.warning('get primary except: %s' %e)
        primary = ''

    return primary

def failover(cluster):
    while True:
        # failover频率不要太快，否则容易导致worker一直返回failed，导致config.databases一直没写进去
        time.sleep(60 + random.randint(0, 30))
        primary = get_primary(cluster)
        if primary == '':
            continue
        mongo_shell_cmd='mongo mongodb://' + cluster + ':' + cluster + '-' + salt_string + \
                        '@' + primary + '/admin --eval "rs.stepDown()"'

        # c1 = 'mongo mongodb://' + cluster + ':' + cluster + '-' + salt_string + '@' + primary + '/db1?authSource=admin --eval "db.t.count()" | egrep -i -v mongodb'
        # c2 = 'mongo mongodb://' + cluster + ':' + cluster + '-' + salt_string + '@' + primary + '/db2?authSource=admin --eval "db.t.count()" | egrep -i -v mongodb'
        # c3 = 'mongo mongodb://' + cluster + ':' + cluster + '-' + salt_string + '@' + primary + '/pytest?authSource=admin --eval "db.t.count()" | egrep -i -v mongodb'
        # print('checking db.t count:', os.popen(c1).read(), os.popen(c2).read(), os.popen(c3).read()) ######################

        # 输出会报： Error: error doing query: failed: network error while attempting to run command 'replSetStepDown'
        # 但实际上触发主下线是成功的，考虑如何避免这一error输出
        res = os.popen(mongo_shell_cmd).read()
        if 'error' in res and "failed: network error while attempting to run command 'replSetStepDown'" not in res:
            failoverlogger.error("FAILURE: stepDown %s in %s, err: %s" %(primary, cluster, res))
        else:
            failoverlogger.info("stepDown %s ok in %s" %(primary, cluster))

def spawn_threads(host, srcport, src, dst, srcpasswd):
    # 消除上次执行在源集群留下的痕迹
    cli = get_client(host, srcport, src, srcpasswd)
    cli.drop_database('pytest')
    # create new user
    try:
        cli.admin.command('createUser', 'pytest1', pwd='pytest', roles=[{'role': 'userAdminAnyDatabase', 'db': 'admin'}])
        cli.pytest.command('createUser', 'pytest2', pwd='pytest', roles=[{'role': 'readWrite', 'db': 'pytest'}])
    except Exception as e:
        logger.error('err: %s' %e)
    cli.close()

    pop = threading.Thread(target=populate, args=(host, srcport, src, srcpasswd))
    # 之所以pop线程不设置成daemon，是想让它在程序结束之前必须停止，否则主线程退出时，它不会退出
    # 因为固定了一个有限的写入量，这样才方便assert 最终的记录数
    pop.start()

    ################################### failover的存在导致CI不是太稳定，后期再启动此线程

    # 主线程退出，failover线程即退出
    fo_src = threading.Thread(target=failover, args=(src,))
    fo_src.setDaemon(True)
    #fo_src.start()

    fo_dst = threading.Thread(target=failover, args=(dst,))
    fo_dst.setDaemon(True)
    #fo_dst.start()

if __name__ ==  '__main__':
    if len(sys.argv) < 2:
        logging.critical("version parameter needed")
        sys.exit(0)

    salt_string = get_salt_string()
    start = time.time()

    version = sys.argv[1]
    # if len(sys.argv) == 3:
    #     # allow to run a single check against one function designated by func string
    #     func=sys.argv[2]

    if get_value_from_env('LOCAL_PYTEST', '') != '':
        local = True

    logger.info("validating cluster version %s through mongos: %s" %(version, MONGOS_HOSTS[0]))

    if version == '3.2':
        spawn_threads(MONGOS_HOSTS[0], 6000, SRC_CLUSTER_3_2, DEST_CLUSTER_3_2, SRC_CLUSTER_3_2 + '-' + salt_string)

        entrypoint(MONGOS_HOSTS[0], 6003, SRC_CLUSTER_3_2, DEST_CLUSTER_3_2, DEST_CLUSTER_3_2 + '-' + salt_string, ETCD_HOST)

    elif version == '3.6':
        spawn_threads(MONGOS_HOSTS[0], 6001, SRC_CLUSTER_3_6, DEST_CLUSTER_3_6, SRC_CLUSTER_3_6 + '-' + salt_string)

        entrypoint(MONGOS_HOSTS[0], 6004, SRC_CLUSTER_3_6, DEST_CLUSTER_3_6, DEST_CLUSTER_3_6 + '-' + salt_string, ETCD_HOST)

    elif version == '4.0':
        spawn_threads(MONGOS_HOSTS[0], 6002, SRC_CLUSTER_4_0, DEST_CLUSTER_4_0, SRC_CLUSTER_4_0 + '-' + salt_string)

        entrypoint(MONGOS_HOSTS[0], 6005, SRC_CLUSTER_4_0, DEST_CLUSTER_4_0, DEST_CLUSTER_4_0 + '-' + salt_string, ETCD_HOST)

    else:
        logger.critical("invalid mongo version")
        sys.exit(-1)

    end = time.time()
    logger.info('pytest end, cluster %s python test took: %f seconds' %(version, (end - start)))

